"""
企业级金融数据MCP服务系统 - 基础数据访问层
提供异步数据访问的基础类和通用功能
"""

from typing import TypeVar, Generic, Type, Optional, List, Dict, Any, Union, Sequence
from datetime import datetime, date
from sqlalchemy.ext.asyncio import AsyncSession
from sqlalchemy import select, insert, update, delete, func, and_, or_, desc, asc
from sqlalchemy.orm import selectinload, joinedload
from sqlalchemy.sql import Select
from sqlalchemy.exc import IntegrityError, NoResultFound
from abc import ABC, abstractmethod
import logging

from src.core.database import BaseModel

# 泛型类型变量
ModelType = TypeVar('ModelType', bound=BaseModel)

logger = logging.getLogger(__name__)


class BaseDAO(Generic[ModelType], ABC):
    """
    基础数据访问对象类
    提供通用的CRUD操作和查询功能
    """
    
    def __init__(self, session: AsyncSession, model_class: Type[ModelType]):
        """
        初始化DAO
        
        Args:
            session: 异步数据库会话
            model_class: 模型类
        """
        self.session = session
        self.model_class = model_class
    
    async def create(self, **kwargs) -> ModelType:
        """
        创建新记录
        
        Args:
            **kwargs: 模型字段值
            
        Returns:
            创建的模型实例
            
        Raises:
            IntegrityError: 违反唯一约束时抛出
        """
        try:
            instance = self.model_class(**kwargs)
            self.session.add(instance)
            await self.session.flush()
            await self.session.refresh(instance)
            return instance
        except IntegrityError as e:
            await self.session.rollback()
            logger.error(f"创建{self.model_class.__name__}记录失败: {e}")
            raise
    
    async def create_many(self, data_list: List[Dict[str, Any]]) -> List[ModelType]:
        """
        批量创建记录
        
        Args:
            data_list: 包含模型字段值的字典列表
            
        Returns:
            创建的模型实例列表
        """
        try:
            instances = [self.model_class(**data) for data in data_list]
            self.session.add_all(instances)
            await self.session.flush()
            for instance in instances:
                await self.session.refresh(instance)
            return instances
        except IntegrityError as e:
            await self.session.rollback()
            logger.error(f"批量创建{self.model_class.__name__}记录失败: {e}")
            raise
    
    async def get_by_id(self, id_value: Any) -> Optional[ModelType]:
        """
        根据ID获取记录
        
        Args:
            id_value: 主键值
            
        Returns:
            模型实例或None
        """
        result = await self.session.execute(
            select(self.model_class).where(self.model_class.id == id_value)
        )
        return result.scalar_one_or_none()
    
    async def get_by_field(self, field_name: str, field_value: Any) -> Optional[ModelType]:
        """
        根据字段值获取记录
        
        Args:
            field_name: 字段名
            field_value: 字段值
            
        Returns:
            模型实例或None
        """
        field = getattr(self.model_class, field_name)
        result = await self.session.execute(
            select(self.model_class).where(field == field_value)
        )
        return result.scalar_one_or_none()
    
    async def get_all(
        self,
        limit: Optional[int] = None,
        offset: Optional[int] = None,
        order_by: Optional[str] = None,
        desc_order: bool = False
    ) -> List[ModelType]:
        """
        获取所有记录
        
        Args:
            limit: 限制返回记录数
            offset: 偏移量
            order_by: 排序字段名
            desc_order: 是否降序排列
            
        Returns:
            模型实例列表
        """
        query = select(self.model_class)
        
        # 添加排序
        if order_by:
            order_field = getattr(self.model_class, order_by)
            if desc_order:
                query = query.order_by(desc(order_field))
            else:
                query = query.order_by(asc(order_field))
        
        # 添加分页
        if offset:
            query = query.offset(offset)
        if limit:
            query = query.limit(limit)
        
        result = await self.session.execute(query)
        return list(result.scalars().all())
    
    async def filter_by(self, **kwargs) -> List[ModelType]:
        """
        根据条件过滤记录
        
        Args:
            **kwargs: 过滤条件
            
        Returns:
            符合条件的模型实例列表
        """
        conditions = []
        for field_name, field_value in kwargs.items():
            if hasattr(self.model_class, field_name):
                field = getattr(self.model_class, field_name)
                conditions.append(field == field_value)
        
        if not conditions:
            return await self.get_all()
        
        query = select(self.model_class).where(and_(*conditions))
        result = await self.session.execute(query)
        return list(result.scalars().all())
    
    async def update_by_id(self, id_value: Any, **kwargs) -> Optional[ModelType]:
        """
        根据ID更新记录
        
        Args:
            id_value: 主键值
            **kwargs: 要更新的字段值
            
        Returns:
            更新后的模型实例或None
        """
        try:
            # 先获取记录
            instance = await self.get_by_id(id_value)
            if not instance:
                return None
            
            # 更新字段
            for field_name, field_value in kwargs.items():
                if hasattr(instance, field_name):
                    setattr(instance, field_name, field_value)
            
            # 更新修改时间
            if hasattr(instance, 'updated_at'):
                instance.updated_at = datetime.utcnow()
            
            await self.session.flush()
            await self.session.refresh(instance)
            return instance
        except IntegrityError as e:
            await self.session.rollback()
            logger.error(f"更新{self.model_class.__name__}记录失败: {e}")
            raise
    
    async def delete_by_id(self, id_value: Any) -> bool:
        """
        根据ID删除记录
        
        Args:
            id_value: 主键值
            
        Returns:
            是否删除成功
        """
        try:
            result = await self.session.execute(
                delete(self.model_class).where(self.model_class.id == id_value)
            )
            return result.rowcount > 0
        except Exception as e:
            await self.session.rollback()
            logger.error(f"删除{self.model_class.__name__}记录失败: {e}")
            raise
    
    async def count(self, **kwargs) -> int:
        """
        统计记录数量
        
        Args:
            **kwargs: 过滤条件
            
        Returns:
            记录数量
        """
        conditions = []
        for field_name, field_value in kwargs.items():
            if hasattr(self.model_class, field_name):
                field = getattr(self.model_class, field_name)
                conditions.append(field == field_value)
        
        query = select(func.count(self.model_class.id))
        if conditions:
            query = query.where(and_(*conditions))
        
        result = await self.session.execute(query)
        return result.scalar() or 0
    
    async def exists(self, **kwargs) -> bool:
        """
        检查记录是否存在
        
        Args:
            **kwargs: 查询条件
            
        Returns:
            是否存在
        """
        count = await self.count(**kwargs)
        return count > 0


class FinancialDataDAO(BaseDAO[ModelType]):
    """
    金融数据专用DAO基类
    提供金融数据特有的查询功能
    """
    
    async def get_by_symbol(self, symbol: str) -> List[ModelType]:
        """
        根据股票代码获取数据
        
        Args:
            symbol: 股票代码
            
        Returns:
            相关数据列表
        """
        if hasattr(self.model_class, 'symbol'):
            return await self.filter_by(symbol=symbol)
        return []
    
    async def get_by_date_range(
        self,
        start_date: Union[date, datetime],
        end_date: Union[date, datetime],
        symbol: Optional[str] = None
    ) -> List[ModelType]:
        """
        根据日期范围获取数据
        
        Args:
            start_date: 开始日期
            end_date: 结束日期
            symbol: 可选的股票代码
            
        Returns:
            指定日期范围内的数据
        """
        conditions = []
        
        # 添加日期条件
        if hasattr(self.model_class, 'trade_date'):
            date_field = self.model_class.trade_date
        elif hasattr(self.model_class, 'date'):
            date_field = self.model_class.date
        elif hasattr(self.model_class, 'created_at'):
            date_field = self.model_class.created_at
        else:
            return []
        
        conditions.append(date_field >= start_date)
        conditions.append(date_field <= end_date)
        
        # 添加股票代码条件
        if symbol and hasattr(self.model_class, 'symbol'):
            conditions.append(self.model_class.symbol == symbol)
        
        query = select(self.model_class).where(and_(*conditions))
        query = query.order_by(desc(date_field))
        
        result = await self.session.execute(query)
        return list(result.scalars().all())
    
    async def get_latest_by_symbol(self, symbol: str, limit: int = 1) -> List[ModelType]:
        """
        获取指定股票的最新数据
        
        Args:
            symbol: 股票代码
            limit: 返回记录数
            
        Returns:
            最新数据列表
        """
        if not hasattr(self.model_class, 'symbol'):
            return []
        
        # 确定日期字段
        date_field = None
        if hasattr(self.model_class, 'trade_date'):
            date_field = self.model_class.trade_date
        elif hasattr(self.model_class, 'date'):
            date_field = self.model_class.date
        elif hasattr(self.model_class, 'created_at'):
            date_field = self.model_class.created_at
        
        query = select(self.model_class).where(self.model_class.symbol == symbol)
        
        if date_field:
            query = query.order_by(desc(date_field))
        
        query = query.limit(limit)
        
        result = await self.session.execute(query)
        return list(result.scalars().all())
    
    async def bulk_upsert(
        self,
        data_list: List[Dict[str, Any]],
        conflict_columns: List[str]
    ) -> int:
        """
        批量插入或更新数据
        
        Args:
            data_list: 数据列表
            conflict_columns: 冲突检测列
            
        Returns:
            影响的行数
        """
        if not data_list:
            return 0
        
        try:
            # 使用PostgreSQL的ON CONFLICT功能
            stmt = insert(self.model_class).values(data_list)
            
            # 构建更新字典（排除冲突列）
            update_dict = {}
            for key in data_list[0].keys():
                if key not in conflict_columns and key != 'id':
                    update_dict[key] = stmt.excluded[key]
            
            # 添加更新时间
            if 'updated_at' in self.model_class.__table__.columns:
                update_dict['updated_at'] = datetime.utcnow()
            
            stmt = stmt.on_conflict_do_update(
                index_elements=conflict_columns,
                set_=update_dict
            )
            
            result = await self.session.execute(stmt)
            return result.rowcount or 0
        except Exception as e:
            await self.session.rollback()
            logger.error(f"批量upsert {self.model_class.__name__}失败: {e}")
            raise


class DAOFactory:
    """DAO工厂类，用于创建各种DAO实例"""
    
    @staticmethod
    def create_dao(session: AsyncSession, model_class: Type[ModelType]) -> BaseDAO[ModelType]:
        """
        创建DAO实例
        
        Args:
            session: 数据库会话
            model_class: 模型类
            
        Returns:
            DAO实例
        """
        # 检查是否是金融数据模型
        if hasattr(model_class, 'symbol') or hasattr(model_class, 'trade_date'):
            return FinancialDataDAO(session, model_class)
        else:
            return BaseDAO(session, model_class)